NVSHMEM 与 JAX/Flax 分布式训练集成实战 · 课程目录

01

NVSHMEM 基础概念

NVSHMEM 是什么 · 与传统 MPI 区别 · 通信模式 (P2P & Collective) · 初始化与终止

02

JAX 分布式基础

pmap / shard_map · Sharded Device Array · 集体通信原语 (psum, pmean, all_gather)

03

Flax 模型定义与分布式

Linen API 简介 · 分布式环境定义 Flax 模型 · 参数分片策略

04

NVSHMEM 内存模型

对称内存分配/释放 · put/get 操作 · 内存 fence 与同步

05

NVSHMEM 点对点通信

nvshmem_put/get · 非阻塞 (nbi) · 信号量 (Signaling) 操作

06

NVSHMEM 集合通信

广播 · 归约/Allreduce · 屏障 · Gather/Allgather 实现

07

JAX 自定义分区与 shard_map

shard_map 语义 · PartitionSpec · 与 pmap 对比

08

JAX 与 NVSHMEM 的桥梁

ctypes/CFFI 调用 C 库 · JIT 嵌入自定义内核 · 数据格式转换

09

Flax 模型参数同步

自定义 Allreduce 梯度同步 · 替代 pmean · 梯度累积与通信重叠

10

流水线并行基础

流水线概念 · 微批次划分 · NVSHMEM 跨节点激活传递

11

张量并行基础

行/列切分 · Transformer Layer 应用 · 通信与计算 Overlap

12

数据并行增强

DDP 瓶颈 · NVSHMEM Allreduce 优化 · Ring Allreduce 实现

13

混合并行策略设计

3D 并行 (数据+张量+流水线) · 通信域划分 · JAX Device Mesh

14

NVSHMEM 与 JAX 异步执行

Async Dispatch · 非阻塞操作 & cudaStream · 回调与同步点

15

性能分析与 Profiling

Nsight Systems · xla_hlo_trace · 识别通信瓶颈

16

错误处理与调试

NVSHMEM 错误码 · NaN/Inf 检测 · pdb 调试 JIT

17

环境搭建与配置

NVIDIA HPC SDK · 环境变量 · CUDA 兼容性 · IB/RoCE 配置

18

案例实战：LLM 预训练 (1)

Flax GPT-2 · NVSHMEM 张量并行 · 数据加载与预处理

19

案例实战：LLM 预训练 (2)

流水线并行 · 梯度检查点 · NVSHMEM 通信优化 · 训练循环

20

案例实战：ViT 图像分类 (1)

Flax Vision Transformer · 数据并行 · 混合精度 AMP

21

案例实战：ViT 图像分类 (2)

Allgather BatchNorm 同步 · 评估/推理通信优化

22

高级主题：NVSHMEM & CUDA Graphs

嵌入 CUDA Graph · JAX AOT 编译 · 性能收益分析

23

高级主题：NVSHMEM & GPUDirect RDMA

零拷贝通信 · RDMA 支持 · JAX 缓冲区注册与传输

24

高级主题：动态形状与通信

动态形状挑战 · 变长数据处理 · Bucketing 策略

25

高级主题：容错与弹性训练

NVSHMEM 容错 · JAX Checkpointing · 通信域重建

26

对比分析：NVSHMEM vs NCCL

延迟/带宽对比 · 编程模型差异 · 场景选择

27

对比分析：NVSHMEM vs MPI

启动方式 · 内存模型 · 集合通信性能 · JAX 集成难易

28

最佳实践与设计模式

通信计算重叠 · 内存池管理 · 避免死锁 · 代码模块化

29

NVSHMEM 2.0 · JAX 自定义通信原语 · 大规模 GPU 趋势

30

综合项目：从零构建分布式训练框架

需求分析 · 架构设计 · 核心模块 (通信/模型/优化器) · 测试部署